import math
from collections import Counter

def knn_predict(train_data, train_labels, test_point, k=3):
    if len(train_data) != len(train_labels):
        raise ValueError("训练数据和标签的数量不匹配")
    if k <= 0:
        raise ValueError("k值必须大于0")

    all_distances_info = []
    for i, point in enumerate(train_data):
        dist = math.sqrt(sum((test_point[j] - point[j])**2 for j in range(len(test_point))))
        all_distances_info.append((dist, point, train_labels[i]))
    
    all_distances_info.sort(key=lambda x: x[0])
    
    k_nearest_neighbors_info = all_distances_info[:k]
    neighbor_labels = [label for _, _, label in k_nearest_neighbors_info]

    vote_counts = Counter(neighbor_labels)
    prediction = vote_counts.most_common(1)[0][0]
        
    return prediction, all_distances_info


if __name__ == "__main__":
    
    print("--- KNN算法详细过程演示 (电影评级预测) ---")
    
    train_features = [
        [95, 5],   
        [88, 2],    
        [110, 15],  
        [150, 200], 
        [105, 8],   
        [140, 150], 
        [98, 10],   
        [160, 250], 
    ]
    train_labels = ['普通', '普通', '大片', '大片', '普通', '大片', '普通', '大片']
    
    test_movie = [120, 50]
    k = 3
    
    print(f"训练集数据 ({len(train_features)} 部):")
    for i, (data, label) in enumerate(zip(train_features, train_labels)):
        print(f"  电影{i+1}: [时长={data[0]}min, 预算={data[1]}M$] -> 评级: {label}")
    
    print(f"\n待预测电影: [时长={test_movie[0]}min, 预算={test_movie[1]}M$]")
    print(f"设定 k={k}")

    prediction, detailed_info = knn_predict(train_features, train_labels, test_movie, k=k)
    
    print("\n--- 详细分析过程 ---")
    print("1. 计算待预测电影与所有已知电影的距离:")
    for i, (dist, point, label) in enumerate(detailed_info):
        print(f"   - 与电影{i+1} ([{point[0]}min, {point[1]}M$], '{label}') 的距离 = {dist:.4f}")

    print(f"\n2. 距离最近的 {k} 个邻居是:")
    k_nearest = detailed_info[:k]
    for i, (dist, point, label) in enumerate(k_nearest):
        print(f"   - 邻居{i+1}: [{point[0]}min, {point[1]}M$], 评级='{label}', 距离={dist:.4f}")

    neighbor_labels_for_vote = [label for _, _, label in k_nearest]
    print(f"\n3. 这 {k} 个邻居的评级分别是: {neighbor_labels_for_vote}")
    
    vote_counts = Counter(neighbor_labels_for_vote)
    print(f"4. 投票统计: {dict(vote_counts)}")
    
    print("\n--- 最终预测结果 ---")
    print(f"待预测电影 (时长{test_movie[0]}分钟, 预算{test_movie[1]}M$) 的预测评级是: '{prediction}'")
